"""
网络拓扑和通信模型

定义了分布式系统的网络结构和通信机制
"""

import numpy as np
import networkx as nx
from typing import Dict, List, Tuple, Optional, Set
import matplotlib.pyplot as plt


class Network:
    """分布式网络类"""
    
    def __init__(self, num_nodes: int, topology: str = 'complete', **kwargs):
        """
        初始化网络
        
        参数:
            num_nodes: 节点数量
            topology: 网络拓扑类型 ('complete', 'ring', 'star', 'random', 'grid')
            **kwargs: 拓扑特定参数
        """
        self.num_nodes = num_nodes
        self.topology = topology
        self.graph = self._create_graph(topology, **kwargs)
        
        # 计算邻接矩阵和拉普拉斯矩阵
        self.adjacency_matrix = nx.adjacency_matrix(self.graph).todense()
        self.laplacian_matrix = nx.laplacian_matrix(self.graph).todense()
        
        # 计算度矩阵
        self.degree_matrix = np.diag(np.array(self.adjacency_matrix.sum(axis=1)).flatten())
        
        # 生成权重矩阵（双随机矩阵）
        self.weight_matrix = self._generate_weight_matrix()
        
    def _create_graph(self, topology: str, **kwargs) -> nx.Graph:
        """创建指定拓扑的图"""
        if topology == 'complete':
            return nx.complete_graph(self.num_nodes)
            
        elif topology == 'ring':
            return nx.cycle_graph(self.num_nodes)
            
        elif topology == 'star':
            return nx.star_graph(self.num_nodes - 1)
            
        elif topology == 'random':
            p = kwargs.get('connectivity', 0.3)
            return nx.erdos_renyi_graph(self.num_nodes, p)
            
        elif topology == 'grid':
            rows = kwargs.get('rows', int(np.sqrt(self.num_nodes)))
            cols = kwargs.get('cols', int(np.ceil(self.num_nodes / rows)))
            G = nx.grid_2d_graph(rows, cols)
            # 重新标记节点为整数
            mapping = {node: i for i, node in enumerate(G.nodes())}
            return nx.relabel_nodes(G, mapping)
            
        else:
            raise ValueError(f"未知的拓扑类型: {topology}")
            
    def _generate_weight_matrix(self) -> np.ndarray:
        """生成双随机权重矩阵"""
        W = np.zeros((self.num_nodes, self.num_nodes))
        
        for i in range(self.num_nodes):
            neighbors = list(self.graph.neighbors(i))
            degree = len(neighbors)
            
            if degree > 0:
                # Metropolis-Hastings权重
                for j in neighbors:
                    degree_j = len(list(self.graph.neighbors(j)))
                    W[i, j] = 1 / (1 + max(degree, degree_j))
                    
                # 对角元素
                W[i, i] = 1 - np.sum(W[i, :])
            else:
                W[i, i] = 1
                
        return W
        
    def get_neighbors(self, node_id: int) -> List[int]:
        """获取节点的邻居列表"""
        return list(self.graph.neighbors(node_id))
        
    def get_degree(self, node_id: int) -> int:
        """获取节点的度"""
        return self.graph.degree(node_id)
        
    def is_connected(self) -> bool:
        """检查网络是否连通"""
        return nx.is_connected(self.graph)
        
    def get_laplacian_eigenvalues(self) -> np.ndarray:
        """获取拉普拉斯矩阵的特征值"""
        eigenvalues = np.linalg.eigvals(self.laplacian_matrix)
        return np.sort(eigenvalues)
        
    def get_weight_eigenvalues(self) -> np.ndarray:
        """获取权重矩阵的特征值"""
        eigenvalues = np.linalg.eigvals(self.weight_matrix)
        return np.sort(eigenvalues)[::-1]
        
    def visualize(self, node_colors: Optional[Dict[int, str]] = None,
                  node_sizes: Optional[Dict[int, float]] = None,
                  title: str = "Network Topology"):
        """可视化网络拓扑"""
        plt.figure(figsize=(10, 8))
        
        # 设置布局
        if self.topology == 'grid':
            pos = nx.spring_layout(self.graph, iterations=50)
        elif self.topology == 'star':
            pos = nx.spring_layout(self.graph, center=[0])
        else:
            pos = nx.spring_layout(self.graph)
            
        # 节点颜色和大小
        if node_colors is None:
            node_colors = ['lightblue'] * self.num_nodes
        else:
            node_colors = [node_colors.get(i, 'lightblue') for i in range(self.num_nodes)]
            
        if node_sizes is None:
            node_sizes = [300] * self.num_nodes
        else:
            node_sizes = [node_sizes.get(i, 300) for i in range(self.num_nodes)]
            
        # 绘制网络
        nx.draw_networkx_nodes(self.graph, pos, node_color=node_colors,
                              node_size=node_sizes, alpha=0.9)
        nx.draw_networkx_labels(self.graph, pos, font_size=12)
        nx.draw_networkx_edges(self.graph, pos, alpha=0.5)
        
        plt.title(title, fontsize=16)
        plt.axis('off')
        plt.tight_layout()
        plt.show()
        
    def get_communication_matrix(self, matrix_type: str = 'laplacian') -> np.ndarray:
        """
        获取通信矩阵
        
        参数:
            matrix_type: 矩阵类型 ('laplacian', 'weight', 'adjacency')
        """
        if matrix_type == 'laplacian':
            return self.laplacian_matrix
        elif matrix_type == 'weight':
            return self.weight_matrix
        elif matrix_type == 'adjacency':
            return self.adjacency_matrix
        else:
            raise ValueError(f"未知的矩阵类型: {matrix_type}")
            
    def consensus_rate(self) -> float:
        """计算共识收敛速率（基于权重矩阵的第二大特征值）"""
        eigenvalues = self.get_weight_eigenvalues()
        # 第二大特征值的绝对值
        return abs(eigenvalues[1]) if len(eigenvalues) > 1 else 0
        
    def algebraic_connectivity(self) -> float:
        """计算代数连通度（拉普拉斯矩阵的第二小特征值）"""
        eigenvalues = self.get_laplacian_eigenvalues()
        return eigenvalues[1] if len(eigenvalues) > 1 else 0
        
    def diameter(self) -> int:
        """计算网络直径"""
        return nx.diameter(self.graph) if self.is_connected() else float('inf')
        
    def average_path_length(self) -> float:
        """计算平均路径长度"""
        if self.is_connected():
            return nx.average_shortest_path_length(self.graph)
        else:
            return float('inf')
            
    def clustering_coefficient(self) -> float:
        """计算聚类系数"""
        return nx.average_clustering(self.graph)
        
    def get_network_statistics(self) -> Dict[str, float]:
        """获取网络统计信息"""
        stats = {
            'num_nodes': self.num_nodes,
            'num_edges': self.graph.number_of_edges(),
            'is_connected': self.is_connected(),
            'diameter': self.diameter(),
            'average_path_length': self.average_path_length(),
            'clustering_coefficient': self.clustering_coefficient(),
            'algebraic_connectivity': self.algebraic_connectivity(),
            'consensus_rate': self.consensus_rate()
        }
        return stats 